package informationretrieval;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;

public class Precision {

	/**
	 * Private constructor to escape instantiation
	 */
	private Precision() {
	}

	/**
	 * Calculate the mean average precision of a list of ranked documents
	 * 
	 * @param rankedResultDocs
	 *            The list of the results for a given query
	 * @param truthRank
	 *            The true positive list of documents ID per a given query
	 * @param falseNegative
	 *            The document ID's we should ignore
	 * @return The average precision result for the given result set
	 */
	public static double calculateMAP(
			List<Integer> rankedResultDocs, Collection<Integer> truthRank,
			Collection<Integer> falseNegative) {
		double result = 0;
		if (falseNegative == null) {
			falseNegative = new ArrayList<Integer>();
		}

		int hitsCounter = 0;
		double averagePrecisionSum = 0;
		int falsePositive = 0;

		for (int i = 0; i < rankedResultDocs.size(); i++) {
			int docId = rankedResultDocs.get(i);
			if (falseNegative.contains(docId)) {
				falsePositive++;
				continue;
			}

			if (!truthRank.contains(docId)) {
				continue;
			}

			hitsCounter++;

			averagePrecisionSum += (double) hitsCounter
					/ (i + 1 - falsePositive);
		}

		if (hitsCounter != 0) {
			result = averagePrecisionSum / hitsCounter;
		}

		return result;
	}

	/**
	 * Calculate the precision at N from a list of documents
	 * 
	 * @param rankedResultDocs
	 *            The list of the results for a given query
	 * @param truthRank
	 *            The true positive list of documents ID per a given query
	 * @param falseNegative
	 *            The document ID's we should ignore
	 * @param cutoffPositions
	 *            the cutoff positions in the list
	 * @return the precision at N for the given data at the different positions
	 *         N
	 */
	public static HashMap<Integer, Double> calculatePrecisionAtN(
			List<Integer> rankedResultDocs, Collection<Integer> truthRank,
			Collection<Integer> falseNegative, int[] cutoffPositions) {

		HashMap<Integer, Double> precisionAtN = new HashMap<Integer, Double>();
		for (int n : cutoffPositions) {
			precisionAtN.put(
					n,
					calculatePrecisionAtN(rankedResultDocs, truthRank,
							falseNegative, n));
		}

		return precisionAtN;
	}

	/**
	 * Calculate the precision at N from a list of documents
	 * 
	 * @param rankedResultDocs
	 *            The list of the results for a given query
	 * @param truthRank
	 *            The true positive list of documents ID per a given query
	 * @param cutOffPosition
	 *            The list cutoff position
	 * @return the precision at N per a given rank list
	 */
	public static double calculatePrecisionAtN(List<Integer> rankedResultDocs,
			Collection<Integer> truthRank, int cutOffPosition) {
		return calculatePrecisionAtN(rankedResultDocs, truthRank,
				new HashSet<Integer>(), cutOffPosition);
	}

	/**
	 * Calculate the precision at N from a list of documents
	 * 
	 * @param rankedResultDocs
	 *            The list of the results for a given query
	 * @param truthRank
	 *            The true positive list of documents ID per a given query
	 * @param falseNegative
	 *            The document ID's we should ignore
	 * @param cutOffPosition
	 *            The list cutoff position
	 * @return the precision at N per a given rank list
	 */
	public static double calculatePrecisionAtN(List<Integer> rankedResultDocs,
			Collection<Integer> truthRank, Collection<Integer> falseNegative,
			int cutOffPosition) {
		double result = (double) calculateHitsAtN(rankedResultDocs, truthRank,
				falseNegative, cutOffPosition) / cutOffPosition;
		return result;
	}

	/**
	 * Calculate the recall at N from a list of documents by a given cutoff
	 * positions
	 * 
	 * @param rankedResultDocs
	 *            The list of the results for a given query
	 * @param truthRank
	 *            The true positive list of documents ID per a given query
	 * @param falseNegative
	 *            The document ID's we should ignore
	 * @param cutOffPositionArray
	 *            An array of cutoff positions
	 * @return The recall at N per a given result set of documents
	 */
	public static HashMap<Integer, Double> calculateRecallAtN(
			List<Integer> rankedResultDocs, Collection<Integer> truthRank,
			Collection<Integer> falseNegative, int[] cutOffPositionArray) {

		HashMap<Integer, Double> recallAtN = new HashMap<Integer, Double>();
		for (int cutOffPosition : cutOffPositionArray)
			recallAtN.put(
					cutOffPosition,
					calculateRecallAtN(rankedResultDocs, truthRank,
							falseNegative, cutOffPosition));

		return recallAtN;
	}

	/**
	 * Calculate the recall at N from a list of documents by a given cutoff
	 * positions
	 * 
	 * @param rankedResultDocs
	 *            The list of the results for a given query
	 * @param truthRank
	 *            The true positive list of documents ID per a given query
	 * @param cutOffPosition
	 *            the cutoff position in the list
	 * @return The recall at N per a given result set of documents
	 */
	public static double calculateRecallAtN(List<Integer> rankedResultDocs,
			Collection<Integer> truthRank, int cutOffPosition) {
		double result = calculateRecallAtN(rankedResultDocs, truthRank,
				new HashSet<Integer>(), cutOffPosition);
		return result;
	}

	/**
	 * Calculate the recall at N from a list of documents by a given cutoff
	 * position
	 * 
	 * @param rankedResultDocs
	 *            The list of the results for a given query
	 * @param truthRank
	 *            The true positive list of documents ID per a given query
	 * @param falseNegative
	 *            The document ID's we should ignore
	 * @param cutOffPosition
	 *            The list cutoff position
	 * @return The recall at N per a given result set of documents
	 */
	public static double calculateRecallAtN(List<Integer> rankedResultDocs,
			Collection<Integer> truthRank, Collection<Integer> falseNegative,
			int cutOffPosition) {
		int hitsAtN = calculateHitsAtN(rankedResultDocs, truthRank,
				falseNegative, cutOffPosition);
		double result = (double) (hitsAtN / truthRank.size());
		return result;
	}

	/**
	 * Calculate the number of hits to the n'th position in a given ranked
	 * result set of documents
	 * 
	 * @param rankedResultDocs
	 *            The list of the results for a given query
	 * @param truthRank
	 *            The true positive list of documents ID per a given query
	 * @param falseNegative
	 *            The document ID's we should ignore
	 * @param cutOffPosition
	 *            The list cutoff position
	 * @return The result for calculating the hits at the cutoff position
	 */
	public static int calculateHitsAtN(List<Integer> rankedResultDocs,
			Collection<Integer> truthRank, Collection<Integer> falseNegative,
			int cutOffPosition) {

		// Method variables
		int hitsCounter = 0;
		int falsePositive = 0;

		if (cutOffPosition < 1) {
			throw new IllegalArgumentException(
					"The cutoff position has to be at least 1");
		}

		for (int i = 0; i < rankedResultDocs.size(); i++) {
			int docId = rankedResultDocs.get(i);
			if (falseNegative.contains(docId)) {
				falsePositive++;
				continue;
			}

			if (!truthRank.contains(docId)) {
				continue;
			}

			if (i < cutOffPosition + falsePositive) {
				hitsCounter++;
			} else {
				break;
			}
		}

		return hitsCounter;
	}
}
